#include <cuda_runtime.h>
#include "conv2d.h"

extern "C" __global__ void implicit_gemm_v1(mykernelParamType param) { 
    int OhOw = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    int n = blockIdx.z;
    if (OhOw >= param.Oh*param.Ow || k >= param.k || n >= param.n) return;

    int oh = OhOw / param.Ow;
    int ow = OhOw % param.Ow;
    float sum = 0.0f;

    for (int crs = 0; crs < param.c*param.r*param.s; crs++) {
        int c = crs / (param.r * param.s);
        int rs = crs % (param.r * param.s);
        int r = rs / param.s;
        int s = rs % param.s;
        
        int ih = oh * param.u - param.p + r;
        int iw = ow * param.v - param.q + s;

        if (ih >= 0 && ih < param.h && iw >= 0 && iw < param.w) {
            size_t input_addr = n * param.c * param.h * param.w 
                               + c * param.h * param.w 
                               + ih * param.w 
                               + iw;
            size_t weight_addr = k * param.c * param.r * param.s 
                               + c * param.r * param.s 
                               + r * param.s 
                               + s;
            sum += param.pin[input_addr] * param.pweight[weight_addr];
        }
    }
    
    size_t output_addr = n * param.k * param.Oh * param.Ow 
                       + k * param.Oh * param.Ow 
                       + oh * param.Ow 
                       + ow;
    param.pout[output_addr] = sum;
}

void launch_implicit_gemm_v1(unsigned int outh, unsigned int outw, unsigned int k, unsigned int n, mykernelParamType* param) {
    dim3 block(16, 16, 1);
    dim3 grid((param->Oh * param->Ow + 15)/16, (param->k + 15)/16, param->n);
    implicit_gemm_v1<<<grid, block>>>(*param);
}